package com.by.datasource.interceptor;

import com.baomidou.mybatisplus.core.toolkit.EncryptUtils;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.parser.JsqlParserSupport;
import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
import com.by.toolkit.exceptions.ServiceException;
import com.google.common.collect.Sets;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.BinaryExpression;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Function;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.expression.operators.relational.NotEqualsTo;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import static com.by.toolkit.utils.ExceptionUtil.nullError;
import static com.by.toolkit.utils.ExceptionUtil.trueError;

@Slf4j
public class BySqlAuditInterceptor extends JsqlParserSupport implements InnerInterceptor {

    /**
     * 缓存验证结果，提高性能
     */
    private static final Set<String> CACHE_VALID_RESULT = Sets.newHashSet();
    /**
     * 缓存表的索引信息
     */
    private static final Map<String, List<IndexInfo>> INDEX_INFO_MAP = new ConcurrentHashMap<>();

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
        MappedStatement ms = mpSh.mappedStatement();
        SqlCommandType sct = ms.getSqlCommandType();
        if (SqlCommandType.INSERT == sct) {
            return;
        }
        BoundSql boundSql = mpSh.boundSql();
        String originalSql = boundSql.getSql();
        String md5Base64 = EncryptUtils.md5Base64(originalSql);
        if (CACHE_VALID_RESULT.contains(md5Base64)) {
            return;
        }
        parserSingle(originalSql, connection);
        //缓存验证结果
        CACHE_VALID_RESULT.add(md5Base64);
    }

    @Override
    protected void processSelect(Select select, int index, String sql, Object obj) {
        SelectBody selectBody = select.getSelectBody();
        if (!(selectBody instanceof PlainSelect)) {
            return;
        }
        PlainSelect plainSelect = (PlainSelect) selectBody;
        Expression where = plainSelect.getWhere();
        //nullError(where, "非法SQL，必须要有where条件");
        Table table = (Table) plainSelect.getFromItem();
        List<Join> joins = plainSelect.getJoins();
        validWhere(where, table, sql, (Connection) obj);
        validJoins(joins, table, sql, (Connection) obj);
    }

    @Override
    protected void processUpdate(Update update, int index, String sql, Object obj) {
        Expression where = update.getWhere();
        nullError(where, "非法SQL，必须要有where条件");
        Table table = update.getTable();
        List<Join> joins = update.getJoins();
        trueError(CollectionUtils.isNotEmpty(joins), "非法SQL，update不准有join");
        validWhere(where, table, sql, (Connection) obj);
        validJoins(joins, table, sql, (Connection) obj);
    }

    @Override
    protected void processDelete(Delete delete, int index, String sql, Object obj) {
        Expression where = delete.getWhere();
        nullError(where, "非法SQL，必须要有where条件");
        Table table = delete.getTable();
        List<Join> joins = delete.getJoins();
        trueError(CollectionUtils.isNotEmpty(joins), "非法SQL，del不准有join");
        validWhere(where, table, sql, (Connection) obj);
        validJoins(joins, table, sql, (Connection) obj);
    }


    /**
     * 验证expression对象是不是 or、not等等
     *
     * @param expression ignore
     */
    private void validExpression(Expression expression, String sql) {
        if (expression instanceof OrExpression) {
            log.error("[BySqlAuditInterceptor][validExpression] finsOrSql:[{}]", sql);
            return;
        }

        if (expression instanceof NotEqualsTo) {
            log.error("[BySqlAuditInterceptor][validExpression] finsNotEqualsSql:[{}]", sql);
            return;
        }

        if (expression instanceof InExpression) {
            InExpression inExpression = (InExpression) expression;
            ItemsList itemsList = inExpression.getRightItemsList();
            if (itemsList instanceof SubSelect) {
                log.warn("[BySqlAuditInterceptor][validExpression] findSubSelectSql:[{}]", sql);
            }
            trueError(itemsList instanceof SubSelect, "非法SQL，能使用子查询");

            if (itemsList instanceof MultiExpressionList) {
                MultiExpressionList multiExpression = (MultiExpressionList) itemsList;
                if (multiExpression.getExpressionLists().size() >= 100) {
                    log.error("[BySqlAuditInterceptor][validExpression] findBigInSql:[{}]", sql);
                }
            }
            return;
        }

        if (expression instanceof BinaryExpression) {
            BinaryExpression binaryExpression = (BinaryExpression) expression;
            if (binaryExpression.getLeftExpression() instanceof Function) {
                Function function = (Function) binaryExpression.getLeftExpression();
                log.warn("[BySqlAuditInterceptor][validExpression] fundFunctionSql:[{}] function:[{}]", sql, function.toString());
            }
            if (binaryExpression.getRightExpression() instanceof SubSelect) {
                log.warn("[BySqlAuditInterceptor][validExpression] findSubSelectSql:[{}]", sql);
                throw new ServiceException("非法SQL，不能使用子查询");
            }
        }

    }

    /**
     * 如果SQL用了 left Join，验证是否有or、not等等，并且验证是否使用了索引
     *
     * @param joins      ignore
     * @param table      ignore
     * @param connection ignore
     */
    private void validJoins(List<Join> joins, Table table, String sql, Connection connection) {
        //允许执行join，验证jion是否使用索引等等
        if (CollectionUtils.isEmpty(joins)) {
            return;
        }
        joins.forEach(join -> {
            Table rightTable = (Table) join.getRightItem();
            join.getOnExpressions().forEach(expression -> {
                validWhere(expression, table, rightTable, sql, connection);
            });
        });
    }

    /**
     * 检查是否使用索引
     *
     * @param table      ignore
     * @param columnName ignore
     * @param connection ignore
     */
    private void validUseIndex(Table table, String columnName, Connection connection) {
        /*//是否使用索引
        boolean useIndexFlag = false;

        String tableInfo = table.getName();
        //表存在的索引
        String dbName = null;
        String tableName;
        String[] tableArray = tableInfo.split("\\.");
        if (tableArray.length == 1) {
            tableName = tableArray[0];
        } else {
            dbName = tableArray[0];
            tableName = tableArray[1];
        }
        List<IndexInfo> indexInfos = getIndexInfos(dbName, tableName, connection);
        for (IndexInfo indexInfo : indexInfos) {
            if (null != columnName && columnName.equalsIgnoreCase(indexInfo.getColumnName())) {
                useIndexFlag = true;
                break;
            }
        }
        if (!useIndexFlag) {
            throw new MybatisPlusException("非法SQL，SQL未使用到索引, table:" + table + ", columnName:" + columnName);
        }*/
    }

    /**
     * 验证where条件的字段，是否有not、or等等，并且where的第一个字段，必须使用索引
     *
     * @param expression ignore
     * @param table      ignore
     * @param connection ignore
     */
    private void validWhere(Expression expression, Table table, String sql, Connection connection) {
        validWhere(expression, table, null, sql, connection);
    }

    /**
     * 验证where条件的字段，是否有not、or等等，并且where的第一个字段，必须使用索引
     *
     * @param expression ignore
     * @param table      ignore
     * @param joinTable  ignore
     * @param connection ignore
     */
    private void validWhere(Expression expression, Table table, Table joinTable, String sql, Connection connection) {
        validExpression(expression, sql);
        if (expression instanceof BinaryExpression) {
            //获得左边表达式
            Expression leftExpression = ((BinaryExpression) expression).getLeftExpression();
            validExpression(leftExpression, sql);

            //如果左边表达式为Column对象，则直接获得列名
            if (leftExpression instanceof Column) {
                Expression rightExpression = ((BinaryExpression) expression).getRightExpression();
                if (joinTable != null && rightExpression instanceof Column) {
                    if (Objects.equals(((Column) rightExpression).getTable().getName(), table.getAlias().getName())) {
                        validUseIndex(table, ((Column) rightExpression).getColumnName(), connection);
                        validUseIndex(joinTable, ((Column) leftExpression).getColumnName(), connection);
                    } else {
                        validUseIndex(joinTable, ((Column) rightExpression).getColumnName(), connection);
                        validUseIndex(table, ((Column) leftExpression).getColumnName(), connection);
                    }
                } else {
                    //获得列名
                    validUseIndex(table, ((Column) leftExpression).getColumnName(), connection);
                }
            }
            //如果BinaryExpression，进行迭代
            else if (leftExpression instanceof BinaryExpression) {
                validWhere(leftExpression, table, joinTable, sql, connection);
            }

            //获得右边表达式，并分解
            Expression rightExpression = ((BinaryExpression) expression).getRightExpression();
            validExpression(rightExpression, sql);
        }
    }

    /**
     * 得到表的索引信息
     *
     * @param dbName    ignore
     * @param tableName ignore
     * @param conn      ignore
     * @return ignore
     */
    public List<IndexInfo> getIndexInfos(String dbName, String tableName, Connection conn) {
        return getIndexInfos(null, dbName, tableName, conn);
    }

    /**
     * 得到表的索引信息
     *
     * @param key       ignore
     * @param dbName    ignore
     * @param tableName ignore
     * @param conn      ignore
     * @return ignore
     */
    public List<IndexInfo> getIndexInfos(String key, String dbName, String tableName, Connection conn) {
        List<IndexInfo> indexInfos = null;
        if (StringUtils.isNotBlank(key)) {
            indexInfos = INDEX_INFO_MAP.get(key);
        }
        if (indexInfos == null || indexInfos.isEmpty()) {
            ResultSet rs;
            try {
                DatabaseMetaData metadata = conn.getMetaData();
                String catalog = StringUtils.isBlank(dbName) ? conn.getCatalog() : dbName;
                String schema = StringUtils.isBlank(dbName) ? conn.getSchema() : dbName;
                rs = metadata.getIndexInfo(catalog, schema, tableName, false, true);
                indexInfos = new ArrayList<>();
                while (rs.next()) {
                    //索引中的列序列号等于1，才有效
                    if (Objects.equals(rs.getString(8), "1")) {
                        IndexInfo indexInfo = new IndexInfo();
                        indexInfo.setDbName(rs.getString(1));
                        indexInfo.setTableName(rs.getString(3));
                        indexInfo.setColumnName(rs.getString(9));
                        indexInfos.add(indexInfo);
                    }
                }
                if (StringUtils.isNotBlank(key)) {
                    INDEX_INFO_MAP.put(key, indexInfos);
                }
            } catch (SQLException e) {
                logger.error(String.format("getIndexInfo fault, with key:%s, dbName:%s, tableName:%s", key, dbName, tableName), e);
            }
        }
        return indexInfos;
    }


    /**
     * 索引对象
     */
    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    private static class IndexInfo {

        private String dbName;

        private String tableName;

        private String columnName;

    }

}
